package com.supertool.tong.merger.util;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.lang.System;

import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.input.CombineFileSplit;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class MergeInputFormat extends FileInputFormat<Comparable,Text>{
    private static Logger LOG = LoggerFactory.getLogger(MergeInputFormat.class);
    private Map<Integer,String> slice2location;

    @Override
    public RecordReader<Comparable, Text> createRecordReader(InputSplit split,
                                                             TaskAttemptContext context) throws IOException,
            InterruptedException {
        return new MergeRecordReader(context.getConfiguration());
    }


    /*
     * Do not split a file into pieces
     * @see org.apache.hadoop.mapreduce.lib.input.FileInputFormat#isSplitable(org.apache.hadoop.mapreduce.JobContext, org.apache.hadoop.fs.Path)
     */
    protected boolean isSplitable(JobContext context, Path filename) {
        return false;
    }

    protected String []getLocationsBySlice(int slice){
        if(slice2location.containsKey(slice)){
            String [] locs=new String[1];
            locs[0]=slice2location.get(slice);
            return locs;
        }
        return new String[0];
    }

    protected int getSlice(String name){
        int beginIndex = name.lastIndexOf("-") + 1;
        int endIndex = name.lastIndexOf(".");
        if (endIndex == -1 || endIndex < beginIndex) {
            return Integer.parseInt(name.substring(name.lastIndexOf("-") + 1));
        } else {
            return Integer.parseInt(name.substring(name.lastIndexOf("-") + 1, name.lastIndexOf(".")));
        }
        //return Integer.parseInt(name.substring(name.lastIndexOf("-")+1), name.lastIndexOf(".")+1);
    }

    private void loadMeta(JobContext job) throws IOException{
        slice2location=new HashMap<Integer,String>();
        Configuration conf=job.getConfiguration();
        String metafile=conf.get("tong.meta.file");
        LOG.info("META="+metafile);
        //metafile="/metafile";
        if (metafile==null){
            return;
        }
        Path metaPath=new Path(metafile);
        FileSystem fileSystem = metaPath.getFileSystem(conf);
        BufferedReader reader=new BufferedReader(
                new InputStreamReader(fileSystem.open(metaPath)),4096);
        String line;
        while((line=reader.readLine())!=null){
            //assume the file format:[slice,host\n]*
            String []items=line.split(",");
            if(items.length<2){
                continue;
            }
            slice2location.put(Integer.parseInt(items[0]), items[1]);
        }
    }

    public String getKeyFromPath(Path path){
        return path.getName();
    }
    public List<InputSplit> getSplits(JobContext job) throws IOException {
        loadMeta(job);
        List<InputSplit> orig_splits=super.getSplits(job);
        Map<String,List<FileSplit>> map=new HashMap<String,List<FileSplit>>();
        for(InputSplit s:orig_splits){
            FileSplit fs=(FileSplit)s;
            String key=getKeyFromPath(fs.getPath());
            LOG.info("FileSplit="+fs.toString());
            if(!map.containsKey(key)){
                map.put(key, new ArrayList<FileSplit>());
            }
            map.get(key).add(fs);
        }
        List<InputSplit> splits=new ArrayList<InputSplit>();
        for(String key:map.keySet()){
            int size=map.get(key).size();
            Path[] paths=new Path[size];
            long[] starts=new long[size];
            long[] lengths=new long[size];
            String[] locations=getLocationsBySlice(getSlice(key));
            int i=0;
            for(FileSplit fs:map.get(key)){
                paths[i]=fs.getPath();
                starts[i]=0;
                lengths[i]=fs.getLength();
                i++;
            }
            splits.add(new CombineFileSplit(paths,starts,lengths,locations));
        }
        LOG.info("Splits total num="+splits.size());
        return splits;
    }
}
